#!/usr/bin/env python3
"""
定时数据收集脚本
"""

import os
import sys
import logging
import schedule
import time
from datetime import datetime
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

# 设置Django环境
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'Didi.settings.dev')

import django
django.setup()

from ai.services.data_crawler import DataCrawler
from ai.services.elasticsearch_service import ElasticsearchService
from ai.services.smart_customer_service import SmartCustomerService

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('scheduled_collection.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

class ScheduledDataCollection:
    """定时数据收集类"""
    
    def __init__(self):
        self.crawler = DataCrawler()
        self.es_service = ElasticsearchService()
        self.smart_service = SmartCustomerService()
        
    def daily_collection(self):
        """每日数据收集任务"""
        logger.info("开始执行每日数据收集任务...")
        start_time = datetime.now()
        
        try:
            # 1. 执行数据爬取
            logger.info("开始数据爬取...")
            crawled_data = self.crawler.crawl_all_sources()
            logger.info(f"数据爬取完成，获取 {len(crawled_data)} 条数据")
            
            # 2. 更新向量数据库
            logger.info("开始更新向量数据库...")
            vector_db_path = self.crawler.create_vector_database(crawled_data)
            if vector_db_path:
                logger.info("向量数据库更新成功")
            else:
                logger.error("向量数据库更新失败")
                return False
            
            # 3. 同步到Elasticsearch
            logger.info("开始同步到Elasticsearch...")
            if self.es_service.reindex_data(crawled_data):
                logger.info("Elasticsearch同步成功")
            else:
                logger.error("Elasticsearch同步失败")
                return False
            

            
            end_time = datetime.now()
            duration = (end_time - start_time).total_seconds()
            logger.info(f"每日数据收集任务完成，耗时: {duration:.2f} 秒")
            return True
            
        except Exception as e:
            logger.error(f"每日数据收集任务失败: {e}")
            return False
    
    def weekly_full_collection(self):
        """每周完整数据收集任务"""
        logger.info("开始执行每周完整数据收集任务...")
        start_time = datetime.now()
        
        try:
            # 1. 清理旧数据
            logger.info("清理旧数据...")
            self.es_service.clear_index()
            
            # 2. 执行完整数据收集
            logger.info("执行完整数据收集...")
            success = self.daily_collection()
            
            if success:
                end_time = datetime.now()
                duration = (end_time - start_time).total_seconds()
                logger.info(f"每周完整数据收集任务完成，耗时: {duration:.2f} 秒")
                return True
            else:
                logger.error("每周完整数据收集任务失败")
                return False
                
        except Exception as e:
            logger.error(f"每周完整数据收集任务失败: {e}")
            return False
    
    def health_check(self):
        """健康检查任务"""
        logger.info("执行系统健康检查...")
        
        try:
            # 检查ES连接
            if self.es_service.es.ping():
                logger.info("Elasticsearch连接正常")
            else:
                logger.error("Elasticsearch连接异常")
                return False
            
            # 检查知识库统计
            stats = self.es_service.get_statistics()
            if stats:
                logger.info(f"知识库统计: {stats}")
            else:
                logger.warning("无法获取知识库统计")
            

            
            logger.info("系统健康检查完成")
            return True
            
        except Exception as e:
            logger.error(f"系统健康检查失败: {e}")
            return False

def main():
    """主函数"""
    logger.info("启动定时数据收集服务...")
    
    # 创建定时任务实例
    scheduler = ScheduledDataCollection()
    
    # 设置定时任务
    # 每日凌晨2点执行数据收集
    schedule.every().day.at("02:00").do(scheduler.daily_collection)
    
    # 每周日凌晨3点执行完整数据收集
    schedule.every().sunday.at("03:00").do(scheduler.weekly_full_collection)
    
    # 每小时执行健康检查
    schedule.every().hour.do(scheduler.health_check)
    
    # 启动时立即执行一次健康检查
    scheduler.health_check()
    
    logger.info("定时任务已设置:")
    logger.info("- 每日凌晨2点: 数据收集")
    logger.info("- 每周日凌晨3点: 完整数据收集")
    logger.info("- 每小时: 健康检查")
    
    # 运行定时任务
    while True:
        try:
            schedule.run_pending()
            time.sleep(60)  # 每分钟检查一次
        except KeyboardInterrupt:
            logger.info("收到中断信号，正在停止服务...")
            break
        except Exception as e:
            logger.error(f"定时任务运行异常: {e}")
            time.sleep(60)

if __name__ == "__main__":
    main() 